package com.zws.cucumber.snapshot;

import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.db.DataRange;
import org.apache.cassandra.db.Keyspace;
import org.apache.cassandra.db.Memtable;
import org.apache.cassandra.db.commitlog.CommitLogPosition;
import org.apache.cassandra.db.filter.ColumnFilter;
import org.apache.cassandra.db.partitions.PartitionUpdate;
import org.apache.cassandra.db.rows.UnfilteredRowIterator;
import org.apache.cassandra.index.transactions.UpdateTransaction;
import org.apache.cassandra.utils.concurrent.OpOrder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicReference;

import static org.apache.cassandra.db.Keyspace.open;

public class ZwsEmbeddedCassSnapshotManager implements ZwsSnapshotManager<Map<String, List<PartitionUpdate>>> {

//    private static final Logger logger = ZwsColorfulLogger.of(LoggerFactory.getLogger(ZwsEmbeddedCassSnapshotManager.class));
    private static final Logger logger = LoggerFactory.getLogger(ZwsEmbeddedCassSnapshotManager.class);

    private ConcurrentMap<String, Map<String, List<PartitionUpdate>>> snapshots = new ConcurrentHashMap<>();

    @Override
    public String name() {
        return "embedded cassandra";
    }

    @Override
    public ConcurrentMap<String, Map<String, List<PartitionUpdate>>> snapshots() {
        return snapshots;
    }

    @Override
    public Logger logger() {
        return logger;
    }

    @Override
    public void snapshot0(String snapshotName, Map<String, Object> args) {
        Map<String, List<PartitionUpdate>> snapshot = new HashMap<>();
        put(snapshotName, snapshot);

        for (ColumnFamilyStore cfs : open(getKeyspace(args)).getColumnFamilyStores()) {
            Memtable memtable = cfs.getTracker().getView().getCurrentMemtable();
            if (!memtable.isClean()) {
                Memtable.MemtableUnfilteredPartitionIterator partitionIterator = memtable.makePartitionIterator(
                        ColumnFilter.all(cfs.metadata),
                        DataRange.allData(cfs.getPartitioner()),
                        false);
                List<PartitionUpdate> partitionUpdateList = new ArrayList<>();
                while (partitionIterator.hasNext()) {
                    UnfilteredRowIterator rowIterator = partitionIterator.next();
                    PartitionUpdate partitionUpdate = PartitionUpdate.fromIterator(rowIterator, ColumnFilter.all(cfs.metadata));
                    partitionUpdateList.add(partitionUpdate);
                }
                snapshot.put(cfs.getTableName(), partitionUpdateList);
            }
        }
    }

    @Override
    public void restore0(String snapshotName, Map<String, Object> args) {
        Map<String, List<PartitionUpdate>> snapshot = get(snapshotName);
        Keyspace keyspace = open(getKeyspace(args));
        for (Map.Entry<String, List<PartitionUpdate>> en : snapshot.entrySet()) {
            String tableName = en.getKey();
            List<PartitionUpdate> partitionUpdateList = en.getValue();
            ColumnFamilyStore cfs = keyspace.getColumnFamilyStore(tableName);
            Memtable copied = new Memtable(new AtomicReference<>(CommitLogPosition.NONE), cfs);;
            cfs.getTracker().reset(copied);
            try (OpOrder.Group group = Keyspace.writeOrder.start()) {
                for (PartitionUpdate partitionUpdate : partitionUpdateList) {
                    cfs.apply(partitionUpdate, UpdateTransaction.NO_OP, group, CommitLogPosition.NONE);
                }
            }
        }
    }

    private String getKeyspace(Map<String, Object> args) {
        return (String) Optional.ofNullable(args)
                .orElseGet(HashMap::new)
                .getOrDefault("keyspace", "com/zws");
    }

}
